#include <ximera/net/socket.hpp>
#include <ximera/assert.hpp>
#include <ximera/system/timer.hpp>
#include <ximera/common/utils.hpp>

// ---------------------------------------------

#if defined(XIMERA_WIN32) || defined(XIMERA_POCKETPC)

#include <winsock2.h>
#include <WS2tcpip.h>

#if defined(XIMERA_PRAGMA_COMMENT_LIB)
    #if defined(XIMERA_WIN32)
        #pragma comment(lib, "Ws2_32.lib")
    #elif defined(BE_POCKETPC)
        #pragma comment(lib, "Ws2.lib")
    #endif
#endif

#elif defined(XIMERA_LINUX)

#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>

#endif

#include <boost/lexical_cast.hpp>

// ---------------------------------------------

namespace ximera { namespace net
{
    namespace
    {
        struct SockAddr
        {
            sockaddr* sa;
            int len;

            IPVersion version;

            sockaddr_in sa4;
            sockaddr_in6 sa6;
        };

        struct SocketHelper
        {
            SocketHelper()
            {
            #if defined(XIMERA_XBOX)
                XNetStartupParams xnsp;
                ZeroMemory(&xnsp, sizeof(xnsp));

                xnsp.cfgSizeOfStruct = sizeof(xnsp);
                xnsp.cfgFlags = XNET_STARTUP_BYPASS_SECURITY;
                
                if (XNetStartup(&xnsp) != 0)
                    XIMERA_THROW(SocketException, L"XNetStartup failed");
            #endif

            #if defined(XIMERA_POCKETPC) || defined(XIMERA_WIN32) || defined(XIMERA_XBOX)
                WSADATA data;
                if (WSAStartup(0x0202, &data) != 0)
                    XIMERA_THROW(SocketException, L"WSAStartup failed!");
            #endif
            }

            ~SocketHelper()
            {
            #if defined(XIMERA_XBOX)
                XNetCleanup();
            #endif

            #if defined(XIMERA_POCKETPC) || defined(XIMERA_WIN32) || defined(XIMERA_XBOX)
                if (WSACleanup() != 0)
                    XIMERA_THROW(SocketException, L"WSACleanup failed!");
            #endif
            }

            Int32 getLastError() const
            {
            #if defined(XIMERA_WIN32)
                return WSAGetLastError();
            #elif defined(XIMERA_LINUX)
                return errno;
            #endif
            }

            String getErrorDesc(Int32 code)
            {
                return boost::lexical_cast<String>(code);
            }

            String getLastErrorStr()
            {
                return getErrorDesc(getLastError());
            }

            SocketHandle createSocket(ProtocolType protocol)
            {
                int type = SOCK_STREAM;
                int proto = IPPROTO_IP;

                switch (protocol)
                {
                case ProtocolType::Udp:
                    type = SOCK_DGRAM;
                    proto = IPPROTO_UDP;
                    break;

                case ProtocolType::Tcp:
                    type = SOCK_STREAM;
                    proto = IPPROTO_TCP;
                    break;
                }

                SocketHandle result = static_cast<SocketHandle>(::socket(PF_INET, type, proto));
                XIMERA_CHECK(result != INVALID_SOCKET,
                    XIMERA_THROW(SocketException, L"Invalid socket!"));

                return result;
            }

            void resolveAddr(IPEndPoint const& endpoint, SockAddr& addr)
            {
                addr.version = endpoint.address().version();

                switch (addr.version)
                {
                case IPVersion::IPv4:                    
                    addr.sa4.sin_family = AF_INET;
                    addr.sa4.sin_addr.s_addr = htonl(endpoint.address().toInt());
                    addr.sa4.sin_port = htons(static_cast<u_short>(endpoint.port()));
                    addr.sa = reinterpret_cast<sockaddr*>(&addr.sa4);
                    addr.len = sizeof(addr.sa4);
                    break;

                case IPVersion::IPv6:
                    // TODO: Finish
                    addr.sa = reinterpret_cast<sockaddr*>(&addr.sa6);
                    addr.len = sizeof(addr.sa6);
                    break;

                default:
                    XIMERA_THROW(SocketException, L"Invalid address!");
                    break;
                }
            }

            void resolveAddr(IPVersion version, SockAddr& addr)
            {
                addr.version = version;

                switch (addr.version)
                {
                case IPVersion::IPv4:
                    addr.sa = reinterpret_cast<sockaddr*>(&addr.sa4);
                    addr.len = sizeof(addr.sa4);
                    break;

                case IPVersion::IPv6:
                    addr.sa = reinterpret_cast<sockaddr*>(&addr.sa6);
                    addr.len = sizeof(addr.sa6);
                    break;

                default:
                    XIMERA_THROW(SocketException, L"Invalid address!");
                    break;
                }
            }

            void resolveAddr(SockAddr& addr, IPEndPoint& endpoint)
            {
                switch (addr.version)
                {
                case IPVersion::IPv4:
                    endpoint.address(IPAddress(string_util::str2wstr(inet_ntoa(addr.sa4.sin_addr))));
                    endpoint.port(addr.sa4.sin_port);
                    break;

                case IPVersion::IPv6:
                    // TODO
                    break;

                default:
                    XIMERA_THROW(SocketException, L"Invalid address!");
                    break;
                }
            }

            UInt32 receive(SocketHandle handle, void* buffer, UInt32 length)
            {
                int result = ::recv(handle, reinterpret_cast<char*>(buffer), length, 0);
                if (result == SOCKET_ERROR)
                {
                    Int32 lastError = getLastError();
                    if (lastError != WSAEWOULDBLOCK)
                        XIMERA_THROW(SocketException, L"Error on receive! (" + getErrorDesc(lastError) + L")");

                    return 0;
                }

                return result;
            }

            UInt32 receive(SocketHandle handle, void* buffer, UInt32 length, SockAddr& addr)
            {
                int result = ::recvfrom(handle, reinterpret_cast<char*>(buffer), length, 0, addr.sa, &addr.len);
                if (result == SOCKET_ERROR)
                {
                    Int32 lastError = getLastError();
                    if (lastError != WSAEWOULDBLOCK)
                        XIMERA_THROW(SocketException, L"Error on receive! (" + getErrorDesc(lastError) + L")");

                    return 0;
                }

                return result;
            }

            UInt32 send(SocketHandle handle, void const* buffer, UInt32 length)
            {
                int result = ::send(handle, reinterpret_cast<char const*>(buffer), length, 0);
                if (result == SOCKET_ERROR)
                {
                    Int32 lastError = getLastError();
                    if (lastError != WSAEWOULDBLOCK)
                        XIMERA_THROW(SocketException, L"Error on send! (" + getErrorDesc(lastError) + L")");
                }

                return result;
            }

            UInt32 send(SocketHandle handle, void const* buffer, UInt32 length, SockAddr const& addr)
            {
                int result = ::sendto(handle, reinterpret_cast<char const*>(buffer), length, 0, addr.sa, addr.len);
                if (result == SOCKET_ERROR)
                {
                    Int32 lastError = getLastError();
                    if (lastError != WSAEWOULDBLOCK)
                        XIMERA_THROW(SocketException, L"Error on send! (" + getErrorDesc(lastError) + L")");
                }

                return result;
            }
        };

        SocketHelper helper;
    }

    Socket::Socket(ProtocolType protocol)
        : _blocking(true), _connected(false), _bound(false), _receiveTimeout(1000)
    {
        _handle = helper.createSocket(protocol);
    }

    Socket::~Socket()
    {
        close();
    }

    bool Socket::blocking() const
    {
        return _blocking;
    }

    void Socket::blocking(bool value)
    {
        if (_blocking == value)
            return;

        u_long arg = value ? 0 : 1;
        int result = ::ioctlsocket(_handle, FIONBIO, &arg);
        XIMERA_CHECK(result == 0,
            XIMERA_THROW(SocketException, L"Couldn't set socket blocking option! (" + helper.getLastErrorStr() + L")"));

        _blocking = value;
    }

    UInt32 Socket::receiveTimeout() const
    {
        return _receiveTimeout;
    }

    void Socket::receiveTimeout(UInt32 value)
    {
        _receiveTimeout = value;
    }

    bool Socket::connected() const
    {
        return _connected;
    }

    bool Socket::bound() const
    {
        return _bound;
    }

    bool Socket::available() const
    {
        fd_set read;
        FD_ZERO(&read);
        FD_SET(_handle, &read);

        timeval timeout;
        timeout.tv_sec = 0;
        timeout.tv_usec = 10000;

        return ::select(1, &read, 0, 0, &timeout) == 1;
    }

    IPEndPoint const& Socket::localEndPoint() const
    {
        return _localEndPoint;
    }

    IPEndPoint const& Socket::remoteEndPoint() const
    {
        return _remoteEndPoint;
    }

    void Socket::connect(IPEndPoint const& endpoint)
    {
        SockAddr addr;
        helper.resolveAddr(endpoint, addr);

        int result = ::connect(_handle, addr.sa, addr.len);
        if (result == SOCKET_ERROR)
        {
            Int32 lastError = helper.getLastError();
            if (_blocking || lastError != WSAEWOULDBLOCK)
                XIMERA_THROW(SocketException, L"Couldn't connect socket! (" + helper.getErrorDesc(lastError) + L")");
        }

        _remoteEndPoint = endpoint;
        _connected = true;
    }

    void Socket::disconnect()
    {
        close();
        // TODO/FIXME: close the connection, but keep object/socket reusable
    }

    void Socket::bind(IPEndPoint const& endpoint)
    {
        SockAddr addr;
        helper.resolveAddr(endpoint, addr);

        int result = ::bind(_handle, addr.sa, addr.len);
        XIMERA_CHECK(result != SOCKET_ERROR,
            XIMERA_THROW(SocketException, L"Couldn't bind socket! (" + helper.getLastErrorStr() + L")"));

        _localEndPoint = endpoint;
        _bound = true;
    }

    void Socket::listen()
    {
        int result = ::listen(_handle, SOMAXCONN);
        XIMERA_CHECK(result != SOCKET_ERROR,
            XIMERA_THROW(SocketException, L"Couldn't listen on socket! (" + helper.getLastErrorStr() + L")"));
    }

    void Socket::accept(Socket& socket)
    {
        SockAddr addr;
        helper.resolveAddr(_localEndPoint.address().version(), addr);

        SocketHandle handle = static_cast<SocketHandle>(::accept(_handle, addr.sa, &addr.len));
        XIMERA_CHECK(handle != INVALID_SOCKET,
            XIMERA_THROW(SocketException, L"Couldn't accept on socket! (" + helper.getLastErrorStr() + L")"));

        // reset socket
        socket._handle = handle;
        helper.resolveAddr(addr, socket._remoteEndPoint);
    }

    void Socket::close()
    {        
        //::shutdown(_handle, 2);

    #if defined(XIMERA_LINUX)
        ::close(_handle);
    #else
        ::closesocket(_handle);
    #endif
    }

    UInt32 Socket::receive(void* buffer, UInt32 length)
    {
        UInt32 received = 0;

        if (_blocking)
        {
            char* bbuffer = reinterpret_cast<char*>(buffer);
            system::Timer timer;
            
            while (received == 0 && (timer.elapsed() * 1000) < _receiveTimeout)
            {
                UInt32 result = helper.receive(_handle, bbuffer, length - received);
                received += result;
            }

            if (received == 0)
                XIMERA_THROW(SocketException, L"Receive timeout was exceeded");
        }
        else
        {
            received = helper.receive(_handle, buffer, length);
            if (received == 0)
                XIMERA_THROW(SocketException, L"No data available!");
        }

        return received;
    }

    UInt32 Socket::receive(void* buffer, UInt32 length, IPEndPoint const& endpoint)
    {
        SockAddr addr;
        helper.resolveAddr(endpoint, addr);

        UInt32 received = 0;

        if (_blocking)
        {
            char* bbuffer = reinterpret_cast<char*>(buffer);
            system::Timer timer;

            while (received == 0 && (timer.elapsed() * 1000) < _receiveTimeout)
            {
                UInt32 result = helper.receive(_handle, bbuffer, length - received, addr);
                received += result;
            }

            if (received == 0)
                XIMERA_THROW(SocketException, L"Receive timeout was exceeded");
        }
        else
        {
            received = helper.receive(_handle, buffer, length, addr);
            if (received == 0)
                XIMERA_THROW(SocketException, L"No data available!");
        }

        return received;
    }

    UInt32 Socket::send(void const* buffer, UInt32 length)
    {
        UInt32 sent = 0;

        if (_blocking)
        {
            char const* bbuffer = reinterpret_cast<char const*>(buffer);            

            while (sent < length)
            {
                UInt32 result = helper.send(_handle, bbuffer, length - sent);

                if (result == 0)
                    XIMERA_THROW(SocketException, L"Connection has been closed");

                sent += result;
                bbuffer += result;
            }
        }
        else
        {
            sent = helper.send(_handle, buffer, length);
        }

        return sent;
    }

    UInt32 Socket::send(void const* buffer, UInt32 length, IPEndPoint const& endpoint)
    {
        SockAddr addr;
        helper.resolveAddr(endpoint, addr);

        UInt32 sent = 0;

        if (_blocking)
        {
            char const* bbuffer = reinterpret_cast<char const*>(buffer);            

            while (sent < length)
            {
                UInt32 result = helper.send(_handle, bbuffer, length - sent, addr);

                if (result == 0)
                    XIMERA_THROW(SocketException, L"Connection has been closed");

                sent += result;
                bbuffer += result;
            }
        }
        else
        {
            sent = helper.send(_handle, buffer, length, addr);
        }

        return sent;
    }
}}
